"""测试各种识别算法的统一接口"""

import numpy as np
import tqdm
import argparse
from utils.ellipse import from_ellipse_args, ellipse_fit
import os

# 识别数据库
from database.plain.binad import XuBuilder, ChenBuilder, ForsythBuilder
from database.plain.triad import HanakBuilder, ParkBuilder, ChristianBuilder
from database.hash import XuHashBuilder, ChristianHashBuilder
from database.hierarchy.binad import ForsythHierarchyBuilder, XuHierarchyBuilder
from database.hierarchy.binad import ChenHierarchyBuilder
from database.hierarchy.triad import ChristianHierarchyBuilder, ParkHierarchyBuilder

# 识别方法
from match.pyramid.binad_fast import XuFastPyramid, ChenFastPyramid, ForsythFastPyramid
from match.pyramid.triad import ParkPyramid, ChristianPyramid
from match.pyramid.triad_vote import ParkPyramidVote, ChristianPyramidVote
from match.hash.binad_vote import XuHashVote
from match.hash.triad import ChristianHash
from match.hash.triad_vote import ChristianHashVote

method_list = {
    # "Xu": (XuBuilder, XuPyramid),  # deprecated
    # "Chen": (ChenBuilder, ChenPyramid), # deprecated
    "binad-plain": {
        "XuFast": (XuBuilder, XuFastPyramid),
        "ChenFast": (ChenBuilder, ChenFastPyramid),
        "ForsythFast": (ForsythBuilder, ForsythFastPyramid),
    },
    "triad-plain": {
        "ParkFast": (ParkBuilder, ParkPyramid),
        "Christian": (ChristianBuilder, ChristianPyramid),
    },
    "hash": {
        "XuHash": (XuHashBuilder, XuHashVote),
        "ChristianHash": (ChristianHashBuilder, ChristianHash),
        "ChristianHashVote": (ChristianHashBuilder, ChristianHashVote),
    },  # # binad hierarchy method
    "binad-hierarchy": {
        "ChenHierarchy": (ChenHierarchyBuilder, ChenFastPyramid),
        "ForsythHierarchy": (ForsythHierarchyBuilder, ForsythFastPyramid),
        "XuHierarchy": (XuHierarchyBuilder, XuFastPyramid),
    },
    # # triad hierarchy method
    "triad-hierarchy": {
        "ChristianHierarchy": (ChristianHierarchyBuilder, ChristianPyramid),
        "ParkHierarchy": (ParkHierarchyBuilder, ParkPyramid),
    },
    # triad vote
    "vote": {
        "ChristianVote": (ChristianBuilder, ChristianPyramidVote),
    },
}


def arg_parse():
    arg_parser = argparse.ArgumentParser()
    arg_parser.add_argument("--noise", type=float, default=0)
    arg_parser.add_argument("--drift", type=float, default=0)
    arg_parser.add_argument(
        "--method",
        type=str,
        default="none",
        choices=[
            "none",
            "binad-plain",
            "triad-plain",
            "hash",
            "binad-hierarchy",
            "triad-hierarchy",
            "vote",
        ],
    )
    arg_parser.add_argument(
        "--exp",
        type=str,
        default="6",
        help="The experiment number in each setting, delimeter by comma",
    )
    arg_parser.add_argument("--outlier", type=float, default=0)
    arg_parser.add_argument("--debug", action="store_true", help="Debug mode")
    arg_parser.add_argument("--catalog", type=str, default="craters.csv")
    arg_parser.add_argument(
        "--exp-files",
        type=str,
        default="./output",
        help="The file to store the experiment",
    )
    arg_parser.add_argument("--device", type=str, default="cuda:0")
    return arg_parser.parse_args()


def debug_triad(module, i, j, k, Qs, C, Qs_prime_idx, idx):
    ix, invariants = module.invariant(i, j, k, Qs[i], Qs[j], Qs[k, None])
    _, real_invariants = module.invariant(
        i, j, k, C[Qs_prime_idx[i]], C[Qs_prime_idx[j]], C[Qs_prime_idx[k], None]
    )
    ind = np.argwhere(
        (idx[0] == Qs_prime_idx[ix[0]])
        & (idx[1] == Qs_prime_idx[ix[1]])
        & (idx[2] == Qs_prime_idx[ix[2]])
    ).squeeze()
    identified_idx = idx[:, ind]
    real_idx = Qs_prime_idx[ix].squeeze()
    return invariants, real_invariants, identified_idx, real_idx


def debug_binad(module, i, j, Qs, C, Qs_prime_idx, idx):
    invariants = module.invariant(Qs[i, None], Qs[j, None])
    real_invariants = module.invariant(
        C[Qs_prime_idx[i], None], C[Qs_prime_idx[j], None]
    )
    ind = np.argwhere(
        (idx[0] == Qs_prime_idx[i]) & (idx[1] == Qs_prime_idx[j])
    ).squeeze()
    identified_idx = idx[:, ind]
    # identified_idx = identifier.idx[:,real_invariants]
    real_idx = Qs_prime_idx[[i, j]].squeeze()
    return invariants, real_invariants, identified_idx, real_idx


def parse_files(exp: str):
    for part in exp.split(","):
        if "-" in part:
            start, end = part.split("-")
            for i in range(int(start), int(end) + 1):
                yield i
        else:
            yield int(part)


if __name__ == "__main__":
    args = arg_parse()
    if args.method == "none":
        print("Please specify the method")
        exit(1)
    method_list = method_list[args.method]
    for k, (BuilderClass, IdentifierClass) in method_list.items():
        name = BuilderClass.__name__
        builder = BuilderClass(
            args.catalog,
            max_field=1000,
            cache_path=f"./__pycache__/{name}.npz",
        )
        catalog_dict = builder()
        catalog = catalog_dict["catalog"]
        descriptor = catalog_dict["descriptor"]
        temp = catalog_dict["temp"]
        idx = catalog_dict["idx"]
        constants = catalog_dict["constants"]
        # 初始化识别器
        identifier = IdentifierClass(f"./__pycache__/{name}.npz", device=args.device)
        method_list[k] = identifier

    # 初始化统计量
    output = {
        name: {"recall": [], "time": [], "failed": 0} for name in method_list.keys()
    }
    # 获取每个测试文件
    exp_files = os.listdir(args.exp_files)
    exp_valid = list(parse_files(args.exp))
    single_folder = os.listdir(f"{args.exp_files}/{exp_files[0]}")
    total = (len(single_folder) - 1) * len(exp_valid)
    gbar = tqdm.tqdm(total=total, desc=f"Testing {args.noise:.1f}/{args.drift:.1f}/{args.outlier:.1f}")
    for exp in exp_valid:
        extra_args = np.loadtxt(f"{args.exp_files}/{exp}/arguments.csv", delimiter=",")
        for file_idx in os.listdir(f"{args.exp_files}/{exp}"):
            if file_idx == "arguments.csv":
                continue
            # 读取世界椭圆和图像椭圆
            data = np.load(f"{args.exp_files}/{exp}/{file_idx}")
            file_idx = int(file_idx.split(".")[0]) - 1
            cam = extra_args[file_idx, :3]
            K = np.array(
                [
                    [cam[0] * 1e3 / cam[1], 0, cam[2] / 2],
                    [0, cam[0] * 1e3 / cam[1], cam[2] / 2],
                    [0, 0, 1],
                ]
            )
            # 加入噪声和漂移
            Qs_prime_idx = data["idx"].astype(np.int32)
            points_2d = data["points"]
            # 加入噪声
            ## 中心点加入噪声和漂移
            points_2d += np.random.normal(args.drift, args.noise, points_2d.shape)
            Qs, Us = ellipse_fit(points_2d[..., 0], points_2d[..., 1])
            # 加入异常值点，即误检测的椭圆，同时删除一个正确的椭圆
            ind = np.random.rand(points_2d.shape[0]) < args.outlier
            if ind.any():
                N = ind.sum()
                r = np.random.uniform(10, 100, (2, N))
                theta = np.random.uniform(0, 2 * np.pi, (1, N))
                E = from_ellipse_args(
                    *np.vstack([np.random.uniform(0, cam[2], (2, N)), r, theta])
                )
                Qs[ind], U = ellipse_fit(E[:, 0], E[:, 1])
                Us[ind] = U
            # 考虑真实情况，同一幅图像中椭圆数量少于4个一定测不出位姿，大于20个不符合真实导航场景
            # # 测试图像平面上的椭圆不变量
            # invariants, real_invariants, identified_idx, real_idx = debug_binad(
            #     identifier, 0, 1, Qs, C, Qs_prime_idx, idx
            # )
            # 对于Park方法的调试
            # invariants, real_invariants, identified_idx, real_idx = debug_triad(
            #     identifier, 0, 1, 2, Qs, C, Qs_prime_idx, idx
            # )
            # 测试识别结果
            for method, identifier in method_list.items():
                t, result = identifier(
                    Qs,
                    Us,
                    factor=0.3,
                    K=K,
                    confidence=0.8,
                    th=3,
                    max_iter=10,
                    max_try=10,
                )
                output[method]["time"].append(t)
                if result is None:
                    output[method]["failed"] += 1
                    continue
                # GT值 : Qs_prime_idx
                # 计算识别准确率
                output[method]["recall"].append(
                    (result[1] == Qs_prime_idx[result[0]]).mean()
                )
            gbar.update(1)
    # 输出结果
    if not os.path.exists("results"):
        os.makedirs("results")
    for method, identifier in method_list.items():
        with open(f"results/{method}.csv", "a") as w:
            w.write(f"{args.noise},{args.drift},{args.outlier},")
            total_recall = np.mean(output[method]["recall"])
            total_time = np.mean(output[method]["time"])
            total_failed = output[method]["failed"]
            w.write(f"{1-total_failed/total:.3f},{total_time:.3f},{total_recall:.3f}\n")
